{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Face Detection using TAO FaceNet\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\"> \n",
    "\n",
    "## Sample predicition using FaceNet\n",
    "\n",
    "<img align=\"center\" src=\"https://docscontent.nvidia.com/dims4/default/123fa56/2147483647/strip/true/crop/2644x1427+0+0/resize/1440x777!/format/webp/quality/90/?url=https%3A%2F%2Fk3-prod-nvidia-docs.s3.amazonaws.com%2Fbrightspot%2Fsphinx%2F00000187-579f-dfeb-adf7-77df80550000%2Ftao%2Ftao-toolkit%2F_images%2Ffacenet.jpg\" width=\"960\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained resnet18 model and train a ResNet-18 FaceNet model on the WIDERFACE dataset\n",
    "* Prune the trained FaceNet model\n",
    "* Retrain the pruned model to recover lost accuracy\n",
    "* Export the pruned model\n",
    "* Run Inference on the trained model\n",
    "* Export the pruned and retrained model to a .onnx file for deployment to DeepStream SDK\n",
    "* Run inference on the exported .onnx model to verify deployment using TensorRT\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `detectnet_v2` model\n",
    "trained on the `WiderFaces` dataset, which you may deploy via [Triton](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps)\n",
    "or [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "### Table of Contents\n",
    "\n",
    "This notebook shows an example use case of Face Detection using FaceNet in the Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables, map drives, and install dependencies](#head-0)\n",
    "1. [Install the TAO Launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2)\n",
    "    1. [Verify and prepare dataset](#head-2-1)\n",
    "    2. [Prepare tfrecords from kitti format dataset](#head-2-2)\n",
    "    3. [Download pre-trained model](#head-2-3)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Prune trained models](#head-6)\n",
    "7. [Retrain pruned models](#head-7)\n",
    "8. [Evaluate retrained model](#head-8)\n",
    "9. [Visualize inferences](#head-9)\n",
    "10. [Deploy](#head-10)\n",
    "    1. [Int8 Optimization](#head-10-1)\n",
    "    2. [Generate TensorRT engine](#head-10-2)\n",
    "11. [Verify Deployed Model](#head-11)\n",
    "    1. [Inference using TensorRT engine](#head-11-1)\n",
    "    2. [Evaluate TensorRT engine](#head-11-2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables, map drives and install dependencies <a class=\"anchor\" id=\"head-0\"></a>\n",
    "When using the purpose-built pretrained models from NGC, please make sure to set the `$KEY` environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.\n",
    "\n",
    "*Note: Please make sure to remove any stray artifacts/files from the `$LOCAL_PROJECT_DIR` paths as mentioned below, that may have been generated from previous experiments. Having checkpoint files etc may interfere with creating a training graph for a new experiment.*\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users' workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/facenet/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/facenet`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "*Note: This notebook currently is by default set up to run training using 1 GPU. To use more GPU's please update the env variable `$NUM_GPUS` accordingly*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up env variables for cleaner command-line commands.\n",
    "import os\n",
    "\n",
    "%env KEY=nvidia_tlt\n",
    "%env NUM_GPUS=1\n",
    "%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/facenet\n",
    "%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/facenet/data\n",
    "%env SPECS_DIR=/workspace/tao-experiments/facenet/specs\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/facenet\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "# The dataset is expected to be present in $LOCAL_PROJECT_DIR/facenet/data, while the results for the steps\n",
    "# in this notebook will be stored at $LOCAL_PROJECT_DIR/facenet\n",
    "# !PLEASE MAKE SURE TO UPDATE THIS PATH!.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tlt-experiments\n",
    "\n",
    "# $PROJECT_DIR is the path to the sample notebook folder and the dependency folder\n",
    "# $PROJECT_DIR/deps should exist for dependency installation\n",
    "%env PROJECT_DIR=/path/to/local/ngc-collaterals/cv/samples/\n",
    "\n",
    "os.environ[\"LOCAL_DATA_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"facenet/data\"\n",
    ")\n",
    "os.environ[\"LOCAL_EXPERIMENT_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"facenet\"\n",
    ")\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"LOCAL_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "\n",
    "# Showing list of specification files.\n",
    "!ls -rlt $LOCAL_SPECS_DIR\n",
    "!ls -rlt $LOCAL_DATASET_SPECS_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below maps the project directory on your local host to a workspace directory in the TAO docker instance, so that the data and the results are mapped from in and out of the docker. For more information please refer to the [launcher instance](https://docs.nvidia.com/tao/tao-toolkit/text/tao_launcher.html) in the user guide.\n",
    "\n",
    "When running this cell on AWS, update the drive_map entry with the dictionary defined below, so that you don't have permission issues when writing data into folders created by the TAO docker.\n",
    "\n",
    "```json\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "            # Mapping the data directory\n",
    "            {\n",
    "                \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "                \"destination\": \"/workspace/tao-experiments\"\n",
    "            },\n",
    "            # Mapping the specs directory.\n",
    "            {\n",
    "                \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "                \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "            },\n",
    "        ],\n",
    "    \"DockerOptions\": {\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid())\n",
    "    }\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "\n",
    "# Define the dictionary with the mapped drives\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ]\n",
    "}\n",
    "\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(drive_map, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install requirement\n",
    "!pip3 install Cython==0.29.36\n",
    "!pip3 install -r $PROJECT_DIR/deps/requirements-pip.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction on this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have set up virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher wheel.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will be using the Wider Face dataset for the tutorial. To find more details, please visit http://shuoyang1213.me/WIDERFACE/. Please download the [training](https://drive.google.com/file/d/0B6eKvaijfFUDQUUwd21EckhUbWs/view?usp=sharing) and [validation](https://drive.google.com/file/d/0B6eKvaijfFUDd3dIRmpvSk8tLUk/view?usp=sharing) set images, and the ground truth [labels](http://mmlab.ie.cuhk.edu.hk/projects/WIDERFace/support/bbx_annotation/wider_face_split.zip), and place the zip files in `$DATA_DOWNLOAD_DIR`. \n",
    "\n",
    "Notes: \n",
    "1. The detection workflow requires dataset to be in kitti format and we will need to convert it.\n",
    "2. The pretrained facenet is trained on grayscale images of height 416 pixels and width 736 pixels. We will convert wider dataset to this format.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1. Verify and prepare dataset <a class=\"anchor\" id=\"head-2-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the dataset is present\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!if [ ! -f $LOCAL_DATA_DIR/WIDER_train.zip ]; then echo 'Train Image zip file not found, please download.'; else echo 'Found Train Image zip file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/WIDER_val.zip ]; then echo 'Validation Image zip file not found, please download.'; else echo 'Found Validation Image zip file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/wider_face_split.zip ]; then echo 'Label zip file not found, please download.'; else echo 'Found Labels zip file.';fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unpack downloaded datasets to $DATA_DOWNLOAD_DIR.\n",
    "!unzip -u $LOCAL_DATA_DIR/WIDER_train.zip -d $LOCAL_DATA_DIR\n",
    "!unzip -u $LOCAL_DATA_DIR/WIDER_val.zip -d $LOCAL_DATA_DIR\n",
    "!unzip -u $LOCAL_DATA_DIR/wider_face_split.zip -d $LOCAL_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify\n",
    "!ls -l $LOCAL_DATA_DIR/WIDER_train/\n",
    "!ls -l $LOCAL_DATA_DIR/WIDER_val/\n",
    "!ls -l $LOCAL_DATA_DIR/wider_face_split/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert wider train dataset to kitti format\n",
    "!python3 convert_wider_to_kitti.py --input_image_dir=$LOCAL_DATA_DIR/WIDER_train/images \\\n",
    "                                   --input_label_file=$LOCAL_DATA_DIR/wider_face_split/wider_face_train_bbx_gt.txt \\\n",
    "                                   --output_dir=$LOCAL_DATA_DIR/training/ \\\n",
    "                                   --image_height=416 --image_width=736 --grayscale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert wider validation dataset to kitti format\n",
    "!python3 convert_wider_to_kitti.py --input_image_dir=$LOCAL_DATA_DIR/WIDER_val/images \\\n",
    "                                   --input_label_file=$LOCAL_DATA_DIR/wider_face_split/wider_face_val_bbx_gt.txt \\\n",
    "                                   --output_dir=$LOCAL_DATA_DIR/validation/ \\\n",
    "                                   --image_height=416 --image_width=736 --grayscale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify\n",
    "import os\n",
    "\n",
    "DATA_DIR = os.environ.get('LOCAL_DATA_DIR')\n",
    "num_training_images = len(os.listdir(os.path.join(DATA_DIR, \"training/images\")))\n",
    "num_training_labels = len(os.listdir(os.path.join(DATA_DIR, \"training/labels\")))\n",
    "num_val_images = len(os.listdir(os.path.join(DATA_DIR, \"validation/images\")))\n",
    "num_val_labels = len(os.listdir(os.path.join(DATA_DIR, \"validation/labels\")))\n",
    "print(\"Number of images in the training set. {}\".format(num_training_images))\n",
    "print(\"Number of labels in the training set. {}\".format(num_training_labels))\n",
    "print(\"Number of images in the validation set. {}\".format(num_val_images))\n",
    "print(\"Number of labels in the validation set. {}\".format(num_val_labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample kitti label.\n",
    "!cat $LOCAL_DATA_DIR/training/labels/30_Surgeons_Surgeons_30_227.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2. Prepare tf records from kitti format dataset <a class=\"anchor\" id=\"head-2-2\"></a>\n",
    "\n",
    "* Update the tfrecords spec file to take in your kitti format dataset\n",
    "* Create the tfrecords using the detectnet_v2 dataset_convert \n",
    "\n",
    "*Note: TfRecords only need to be generated once.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TFrecords conversion spec file for kitti training\")\n",
    "!cat $LOCAL_SPECS_DIR/facenet_tfrecords_kitti_train.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a new directory for the output tfrecords dump.\n",
    "print(\"Converting Tfrecords for wider train dataset\")\n",
    "!mkdir -p $LOCAL_DATA_DIR/tfrecords && rm -rf $LOCAL_DATA_DIR/tfrecords/*\n",
    "!tao model detectnet_v2 dataset_convert \\\n",
    "                  -d $SPECS_DIR/facenet_tfrecords_kitti_train.txt \\\n",
    "                  -o $DATA_DOWNLOAD_DIR/tfrecords/training/kitti_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a new directory for the output tfrecords dump.\n",
    "print(\"Converting Tfrecords for wider validation dataset\")\n",
    "!tao model detectnet_v2 dataset_convert \\\n",
    "                  -d $SPECS_DIR/facenet_tfrecords_kitti_val.txt \\\n",
    "                  -o $DATA_DOWNLOAD_DIR/tfrecords/validation/kitti_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_DATA_DIR/tfrecords/training/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_DATA_DIR/tfrecords/validation/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3. Download pre-trained model <a class=\"anchor\" id=\"head-2-3\"></a>\n",
    "\n",
    "Download the correct pretrained model from the NGC model registry for your experiment. Please note that for DetectNet_v2, the input is expected to be 0-1 normalized.\n",
    "\n",
    "For FaceNet pretrain model please download from: `nvidia/tao/facenet`\n",
    "\n",
    "After downloading the pre-trained model, please place the files in $LOCAL_EXPERIMENT_DIR\n",
    "You will then have the following path\n",
    "\n",
    "* pretrain model in `$LOCAL_EXPERIMENT_DIR/pretrain_models/facenet_vunpruned_v2.0/model.tlt`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List models available in the model registry.\n",
    "!ngc registry model list nvidia/tao/facenet:unpruned*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Position the pretrain model to the target destination.\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/pretrain_models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/facenet:unpruned_v2.0 \\\n",
    "    --dest $LOCAL_EXPERIMENT_DIR/pretrain_models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the pretrained model is present\n",
    "!if [ ! -f $LOCAL_EXPERIMENT_DIR/pretrain_models/facenet_vunpruned_v2.0/model.tlt ]; then echo 'Pretrain model file not found, please download.'; else echo 'Found Pretrain model file.';fi"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Tfrecords for the train datasets: In order to use the newly generated tfrecords, update the dataset_config parameter in the spec file at `$LOCAL_SPECS_DIR/facenet_train_resnet18_kitti.txt` \n",
    "* Augmentation parameters for on the fly data augmentation\n",
    "* Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc.\n",
    "\n",
    "*Note: Please set the `load_graph` option to `true` in the model_config to load the pretrained facenet model.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/facenet_train_resnet18_kitti.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "\n",
    "*Note: The training may take hours to complete. Also, the remaining notebook, assumes that the training was done in single-GPU mode. When run in multi-GPU mode, please expect to update the pruning and inference steps with new pruning thresholds and updated parameters in the clusterfile.json accordingly for optimum performance.*\n",
    "\n",
    "*Detectnet_v2 now supports restart from checkpoint. In case, the training job is killed prematurely, you may resume training from the closest checkpoint by simply re-running the **same** command line. Please do make sure to use the <u>**same number of GPUs**</u> when restarting the training.*\n",
    "\n",
    "*When running the training with NUM_GPUs>1, you may need to modify the `batch_size_per_gpu` and `learning_rate` to get similar mAP as a 1GPU training run. In most cases, scaling down the batch-size by a factor of NUM_GPU's or scaling up the learning rate by a factor of NUM_GPU's would be a good place to start.* "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we evaluate the pretrained Face Detect network on Wider validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model detectnet_v2 evaluate -e $SPECS_DIR/facenet_train_resnet18_kitti.txt \\\n",
    "                           -k $KEY \\\n",
    "                           -m $USER_EXPERIMENT_DIR/pretrain_models/facenet_vunpruned_v2.0/model.tlt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we train model on Wider train set starting with the pretrained Face Detect model weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model detectnet_v2 train -e $SPECS_DIR/facenet_train_resnet18_kitti.txt \\\n",
    "                        -r $USER_EXPERIMENT_DIR/experiment_dir_unpruned \\\n",
    "                        -k $KEY \\\n",
    "                        -n resnet18_detector \\\n",
    "                        --gpus $NUM_GPUS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('Model for each epoch:')\n",
    "print('---------------------')\n",
    "!ls -lrthR $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate the trained model <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model detectnet_v2 evaluate -e $SPECS_DIR/facenet_train_resnet18_kitti.txt\\\n",
    "                           -k $KEY \\\n",
    "                           -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/resnet18_detector.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Prune the trained model <a class=\"anchor\" id=\"head-6\"></a>\n",
    "* Specify pre-trained model\n",
    "* Equalization criterion\n",
    "* Threshold for pruning.\n",
    "* Output directory to store the model\n",
    "\n",
    "*Usually, you just need to adjust `-pth` (threshold) for accuracy and model size trade off. Higher `pth` gives you smaller model (and thus higher inference speed) but worse accuracy. The threshold to use depends on the dataset. A pth value `5.2e-6` is just a start point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy.*\n",
    "\n",
    "*For some internal studies, we have noticed that a pth value of 0.01 is a good starting point for detectnet_v2 models.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an output directory if it doesn't exist.\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model detectnet_v2 prune \\\n",
    "                  -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/resnet18_detector.hdf5 \\\n",
    "                  -k $KEY \\\n",
    "                  -o $USER_EXPERIMENT_DIR/experiment_dir_pruned/resnet18_nopool_bn_detectnet_v2_pruned.hdf5 \\\n",
    "                  -eq union \\\n",
    "                  -pth 0.0000052"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Retrain the pruned model <a class=\"anchor\" id=\"head-7\"></a>\n",
    "* Model needs to be re-trained to bring back accuracy after pruning\n",
    "* Specify re-training specification with pretrained weights as pruned model.\n",
    "\n",
    "*Note: For retraining, please set the `load_graph` option to `true` in the model_config to load the pruned model graph. Also, if after retraining, the model shows some decrease in mAP, it could be that the originally trained model, was pruned a little too much. Please try reducing the pruning threshold, thereby reducing the pruning ratio, and use the new model to retrain.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Printing the retrain experiment file. \n",
    "# Note: We have updated the experiment file to include the \n",
    "# newly pruned model as a pretrained weights and, the\n",
    "# load_graph option is set to true \n",
    "!cat $LOCAL_SPECS_DIR/facenet_retrain_resnet18_kitti.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retraining using the pruned model as pretrained weights \n",
    "!tao model detectnet_v2 train -e $SPECS_DIR/facenet_retrain_resnet18_kitti.txt \\\n",
    "                        -r $USER_EXPERIMENT_DIR/experiment_dir_retrain \\\n",
    "                        -k $KEY \\\n",
    "                        -n resnet18_detector_pruned \\\n",
    "                        --gpus $NUM_GPUS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listing the newly retrained model.\n",
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain/weights"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Evaluate the retrained model <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This section evaluates the pruned and retrained model, using `tao-evaluate`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model detectnet_v2 evaluate -e $SPECS_DIR/facenet_retrain_resnet18_kitti.txt \\\n",
    "                           -k $KEY \\\n",
    "                           -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/resnet18_detector_pruned.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Visualize inferences <a class=\"anchor\" id=\"head-9\"></a>\n",
    "In this section, we run the `inference` tool to generate inferences on the trained models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running inference for detection on n images\n",
    "!tao model detectnet_v2 inference -e $SPECS_DIR/facenet_inference_kitti_tlt.txt \\\n",
    "                            -r $USER_EXPERIMENT_DIR/tlt_infer_testing \\\n",
    "                            -k $KEY \\\n",
    "                            -i $DATA_DOWNLOAD_DIR/validation/images"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `infer` tool produces two outputs. \n",
    "1. Overlain images in `$USER_EXPERIMENT_DIR/tlt_infer_testing/images_annotated`\n",
    "2. Frame by frame bbox labels in kitti format located in `$USER_EXPERIMENT_DIR/tlt_infer_testing/labels`\n",
    "\n",
    "*Note: To run inferences for a single image, simply replace the path to the -i flag in `infer` command with the path to the image.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg', '.png', '.jpeg', '.ppm']\n",
    "\n",
    "def visualize_images(image_dir, num_cols=4, num_images=10):\n",
    "    output_path = os.path.join(os.environ['LOCAL_EXPERIMENT_DIR'], image_dir)\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Visualizing the first 12 images.\n",
    "OUTPUT_PATH = 'tlt_infer_testing/images_annotated' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 4 # number of columns in the visualizer grid.\n",
    "IMAGES = 12 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls $LOCAL_EXPERIMENT_DIR/data/training"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Deploy! <a class=\"anchor\" id=\"head-10\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_final\n",
    "# Removing a pre-existing copy of the onnx if there has been any.\n",
    "import os\n",
    "output_file=os.path.join(os.environ['LOCAL_EXPERIMENT_DIR'],\n",
    "                         \"experiment_dir_final/resnet18_detector.onnx\")\n",
    "if os.path.exists(output_file):\n",
    "    os.system(\"rm {}\".format(output_file))\n",
    "!tao model detectnet_v2 export \\\n",
    "            -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/resnet18_detector_pruned.hdf5 \\\n",
    "            -k $KEY \\\n",
    "            -o $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector.onnx"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `tao deploy` container, you can generate a TensorRT engine and verify the correctness of the generated through evaluate and inference. \n",
    "\n",
    "The `tao deploy` produces optimized tensorrt engines for the platform that it resides on. Therefore, to get maximum performance, please run `tao deploy` command which will instantiate a deploy container, with the exported `.onnx` file on your target device. The `tao deploy` container only works for x86, with discrete NVIDIA GPU's. \n",
    "\n",
    "For the jetson devices, please download the tao-converter for jetson from the dev zone link [here](https://developer.nvidia.com/tao-converter). \n",
    "\n",
    "If you choose to integrate your model into deepstream directly, you may do so by simply copying the exported `.onnx` file along with the calibration cache to the target device and updating the spec file that configures the `gst-nvinfer` element to point to this newly exported model. Usually this file is called `config_infer_primary.txt` for detection models and `config_infer_secondary_*.txt` for classification models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls $LOCAL_EXPERIMENT_DIR/data/training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Need to pass the actual image directory instead of data root for tao deploy to locate images for calibration\n",
    "!sed -i \"s|/workspace/tao-experiments/facenet/data/training|/workspace/tao-experiments/facenet/data/training/images|g\" $LOCAL_SPECS_DIR/facenet_retrain_resnet18_kitti.txt\n",
    "# Convert to TensorRT engine (INT8)\n",
    "!tao deploy detectnet_v2 gen_trt_engine \\\n",
    "                  -m $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector.onnx \\\n",
    "                  -k $KEY \\\n",
    "                  --data_type int8 \\\n",
    "                  --batches 10 \\\n",
    "                  --batch_size 32 \\\n",
    "                  --max_batch_size 32 \\\n",
    "                  --engine_file $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector.trt.int8 \\\n",
    "                  --cal_cache_file $USER_EXPERIMENT_DIR/experiment_dir_final/calibration.bin \\\n",
    "                  -e $SPECS_DIR/facenet_retrain_resnet18_kitti.txt \\\n",
    "                  --results_dir $USER_EXPERIMENT_DIR/export \\\n",
    "                  --verbose\n",
    "# Convert back the spec file\n",
    "!sed -i \"s|/workspace/tao-experiments/facenet/data/training/images|/workspace/tao-experiments/facenet/data/training|g\" $LOCAL_SPECS_DIR/facenet_retrain_resnet18_kitti.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Note: For this example, we generate a calibration tensorfile containing 10 batches of training data.\n",
    "Ideally, it is best to use at least 10-20% of the training data to do so. The more data provided during calibration, the closer int8 inferences are to fp32 inferences.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lh $LOCAL_EXPERIMENT_DIR/experiment_dir_final"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Verify Deployed Model <a class=\"anchor\" id=\"head-11\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 11.1. Inference using TensorRT engine <a class=\"anchor\" id=\"head-11-1\"></a>\n",
    "\n",
    "Verify the exported model by visualizing inferences on TensorRT.\n",
    "In addition to running inference on a `.hdf5` model in [step 9](#head-9), the `inference` tool is also capable of consuming the converted `TensorRT engine` from [step 10](#head-10)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy detectnet_v2 inference -e $SPECS_DIR/facenet_inference_kitti_etlt.txt \\\n",
    "                                   -m $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector.trt.int8 \\\n",
    "                                   -r $USER_EXPERIMENT_DIR/etlt_infer_testing \\\n",
    "                                   -i $DATA_DOWNLOAD_DIR/validation/images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the first 12 inferenced images.\n",
    "OUTPUT_PATH = 'etlt_infer_testing/images_annotated' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 4 # number of columns in the visualizer grid.\n",
    "IMAGES = 12 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 11.2. Evaluate using TensorRT engine <a class=\"anchor\" id=\"head-11-2\"></a>\n",
    "\n",
    "Verify the exported model by evaluating the TensorRT.\n",
    "In addition to running inference on a `.hdf5` model in [step 9](#head-9), the `inference` tool is also capable of consuming the converted `TensorRT engine` from [step 10](#head-10)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "val_split = 0.14  # Val split from the spec file\n",
    "img_ext = '.png'\n",
    "lab_ext = '.txt'\n",
    "\n",
    "# Get images and labels from the partitioned validation set\n",
    "images_root = f\"{os.environ['LOCAL_DATA_DIR']}/training/images\"\n",
    "labels_root = f\"{os.environ['LOCAL_DATA_DIR']}/training/labels\"\n",
    "images_list = [os.path.splitext(imfile)[0] for imfile in sorted(os.listdir(images_root)) if imfile.endswith(img_ext)]\n",
    "num_val_images = int(len(images_list) * val_split)\n",
    "\n",
    "# Copy the data to a separate directory for evaluation\n",
    "os.makedirs(os.path.join(f\"{os.environ['LOCAL_DATA_DIR']}/training\", \"images_val\"), exist_ok=True)\n",
    "os.makedirs(os.path.join(f\"{os.environ['LOCAL_DATA_DIR']}/training\", \"labels_val\"), exist_ok=True)\n",
    "\n",
    "for fname in images_list[:num_val_images]:\n",
    "    shutil.copyfile(os.path.join(images_root, fname + img_ext), \\\n",
    "                             os.path.join(f\"{os.environ['LOCAL_DATA_DIR']}/training\", \"images_val\", fname + img_ext))\n",
    "    shutil.copyfile(os.path.join(labels_root, fname + lab_ext), \\\n",
    "                             os.path.join(f\"{os.environ['LOCAL_DATA_DIR']}/training\", \"labels_val\", fname + lab_ext))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy detectnet_v2 evaluate -e $SPECS_DIR/facenet_retrain_resnet18_kitti.txt \\\n",
    "                                  -m $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector.trt.int8 \\\n",
    "                                  -i $DATA_DOWNLOAD_DIR/training/images_val \\\n",
    "                                  -l $DATA_DOWNLOAD_DIR/training/labels_val \\\n",
    "                                  -r $USER_EXPERIMENT_DIR/experiment_dir_retrain/"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
